import os
import cv2
import math
import time
import torch
import numpy as np
import torch.nn as nn

#from backbone.resnet_dcn import ResNet
from backbone.dlanet_dcn import DlaNet as DlaNet
# from backbone.dlanet import DlaNet as DlaNet

from loss.Loss import _gather_feat,_transpose_and_gather_feat
from PIL import Image, ImageDraw
from datasets_loader.dataloader_hrsc import get_affine_transform
#from Loss import _transpose_and_gather_feat

from datasets.dataset_hrsc import HRSC
from datasets.calMap import DecDecoder,write_results
import argparse




def parse_args():
    #argparse是python用于解析命令行参数和选项的标准模块
    #很多时候，需要用到解析命令行参数的程序，目的是在终端窗口(ubuntu是终端窗口，windows是命令行窗口)输入训练的参数和选项。
    #可以让人轻松编写用户友好的命令行接口
    parser = argparse.ArgumentParser(description='BBAVectors Implementation')#创建解析器

    parser.add_argument('--input_h', type=int, default=608, help='Resized image height')
    parser.add_argument('--input_w', type=int, default=608, help='Resized image width')
    parser.add_argument('--CSL', type=int, default=1)
    parser.add_argument('--SUM', type=int, default=1)
#     parser.add_argument('--method', default='ce_608_gauss_')#ce_wh_608_gauss  baseline_ratio_
    parser.add_argument('--path',default='./save_weights/HRSC2016_weight/ce_608_gauss_SUM')#ce_gauss_608_r20
    parser.add_argument('--data_dir',default="../BBA/datasets/HRSC/FullDataSet/")
    parser.add_argument('--ovthresh', default=[0.5,0.75])
    args = parser.parse_args()#解析参数
    return args


args = parse_args()


if args.CSL and not args.SUM:
    print("using CSL")
    head = {'hm': 1, 'wh': 2, 'ang':180, 'reg': 2}
elif args.SUM and args.CSL:
    print("using CSL and ang")
    head = {'hm': 1, 'wh': 2, 'ang_csl':180,'reg': 2,'ang':1, }
else:
    print("using baseline")
    head = {'hm': 1, 'wh': 2, 'ang':1, 'reg': 2}


model = DlaNet(34,heads = head )
device = torch.device('cuda')
model.load_state_dict(torch.load(os.path.join(args.path,'best.pth')))#or best.pth?
model.eval()
model.cuda()
down_ratio = 4
decoder = DecDecoder(100,0.1,1)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

data_dir = args.data_dir

dset = HRSC(data_dir=data_dir,input_h=args.input_h, input_w=args.input_w)

def dec_eval(args,model,dsets_,down_ratio,device,decoder):  
    result_path = args.path
    mAP= []
    for ovthresh in args.ovthresh:
        
        write_results(args=args,
                  model=model,
                  dsets=dsets_,
                  down_ratio=down_ratio,
                  device=device,
                  decoder=decoder,
                  result_path=result_path,
                  print_ps=False,)
        map_s,map_m,map_l,map_ = dsets_.dec_evaluation(result_path,ovthresh)       
        mAP.append(["%.4f"%map_s,"%.4f"%map_m,"%.4f"%map_l,"%.4f"%map_])

    return mAP

ap = dec_eval(args,model,dset,down_ratio,device,decoder)
print(ap)